# 1D Hermite程序
import numpy as np
from scipy.integrate import quad
from Hermite import *


# 定义 Hermite 基函数及其导数
def phi_0(x):
    return 1 - 3 * x**2 + 2 * x**3


def phi_1(x):
    return x - 2 * x**2 + x**3


def phi_2(x):
    return 3 * x**2 - 2 * x**3


def phi_3(x):
    return -x**2 + x**3


def dphi_0(x):
    return -6 * x + 6 * x**2


def dphi_1(x):
    return 1 - 4 * x + 3 * x**2


def dphi_2(x):
    return 6 * x - 6 * x**2


def dphi_3(x):
    return -2 * x + 3 * x**2


class Mesh:
    def __init__(self, n: int, interval: list = [0, 1]):
        """n为剖分设置
            interval为区间
        Args:
            n (int): _description_
        """
        mesh = np.linspace(interval[0], interval[1], n+1)
        self.mesh = mesh[:]
        self.h = (interval[1] - interval[0]) / n  # 步长
        self.n = n  # 网格剖分
        self.n_nodes = n + 1  # 节点数
        
        
    

class FE:
    def __init__(self, mesh: Mesh, order: int=3):
        """Hermite有限元空间

        Args:
            mesh (Mesh): 网格
            order (int): 有限元空间的阶数
        """
        self.mesh = mesh  # 网格
        self.n = mesh.n  # 网格剖分
        self.h = mesh.h  # 网格步长
        self.n_nodes = mesh.n_nodes  # 节点数
        self.n_dofs_per_node = 2
        self.n_dofs = self.n_dofs_per_node * mesh.n_nodes  # 自由度个数=节点数*2
        # 初始化刚度矩阵 和 载荷向量
        self.K = np.zeros((self.n_dofs, self.n_dofs))
        self.b = np.zeros(self.n_dofs)
        
    def assemble_matrix(self):
        """
        组装刚度矩阵
        """
        n_dofs_per_element = self.n_dofs_per_node * 2  # 每个单元的自由度数量
        for i in range(self.n):
            # 单元刚度矩阵
            Ke = np.zeros((n_dofs_per_element, n_dofs_per_element))
            for alpha in range(n_dofs_per_element):
                for beta in range(n_dofs_per_element):
                    if alpha == 0:
                        dphi_alpha = dphi_0
                    elif alpha == 1:
                        dphi_alpha = dphi_1
                    elif alpha == 2:
                        dphi_alpha = dphi_2
                    else:
                        dphi_alpha = dphi_3

                    if beta == 0:
                        dphi_beta = dphi_0
                    elif beta == 1:
                        dphi_beta = dphi_1
                    elif beta == 2:
                        dphi_beta = dphi_2
                    else:
                        dphi_beta = dphi_3

                    def integrand(x):
                        return (1 / self.h**2) * dphi_alpha(x) * dphi_beta(x)

                    Ke[alpha, beta] = quad(integrand, 0, 1)[0]

            # 局部到全局的映射
            dof_indices = [self.n_dofs_per_node * i +
                           j for j in range(n_dofs_per_element)]
            for alpha in range(n_dofs_per_element):
                for beta in range(n_dofs_per_element):
                    self.K[dof_indices[alpha],
                           dof_indices[beta]] += Ke[alpha, beta]

    def assemble_vector(self):
        """
        组装载荷向量
        """
        n_dofs_per_element = self.n_dofs_per_node * 2  # 每个单元的自由度数量
        # 定义源项函数

        def f(x):
            return 1.0  # 这里简单假设源项为常数1

        for i in range(self.n):
            # 单元载荷向量
            be = np.zeros(n_dofs_per_element)
            for alpha in range(n_dofs_per_element):
                if alpha == 0:
                    phi_alpha = phi_0
                elif alpha == 1:
                    phi_alpha = phi_1
                elif alpha == 2:
                    phi_alpha = phi_2
                else:
                    phi_alpha = phi_3

                def integrand(x):
                    return self.h * f(self.mesh[i] + x * self.h) * phi_alpha(x)

                be[alpha] = quad(integrand, 0, 1)[0]

            # 局部到全局的映射
            dof_indices = [self.n_dofs_per_node * i +
                           j for j in range(n_dofs_per_element)]
            for alpha in range(n_dofs_per_element):
                self.b[dof_indices[alpha]] += be[alpha]


# test
if __name__ == "__main__":
    mesh = Mesh(8)
    print(mesh.h)
    fe = FE(mesh, 3)
    fe.assemble_matrix()
    fe.assemble_vector()
    print("刚度矩阵 K:")
    print(fe.K)
    print("载荷向量 b:")
    print(fe.b)
